import networkx as nx
import matplotlib.pyplot as plt
from networkx.algorithms.community import kernighan_lin_bisection


# 读取数据集
G = nx.read_edgelist('data/email-Network.txt', create_using=nx.Graph())
# department_labels因为只有一个‘部门’属性用字典的方式存储
department_labels = {}
with open('data/email-Department.txt', 'r') as file:
    for line in file:
        # split()默认以空格的方式分割
        node, label = line.strip().split()
        department_labels[node] = label
# 将部门标签信息添加到节点属性中
nx.set_node_attributes(G, department_labels, 'department')

def draw_spring(G, com):
    """
    G:图
    com：划分好的社区
    node_size表示节点大小
    node_color表示节点颜色
    node_shape表示节点形状
    with_labels=True表示节点是否带标签
    """
    pos = nx.spring_layout(G)  # 节点的布局为spring型
    NodeId = list(G.nodes())
    node_size = [G.degree(i) ** 1.2 * 90 for i in NodeId]  # 节点大小

    plt.figure(figsize=(8, 6))  # 图片大小
    nx.draw(G, pos, with_labels=True, node_size=node_size, node_color='w', node_shape='.')

    color_list = ['pink', 'orange', 'r', 'g', 'b', 'y', 'm', 'gray', 'black', 'c', 'brown',

                '#FF48FF',
                '#DDE1E4',
                '#52FF7E',
                '#FF609D',
                '#C6B2E6',
                '#1EA8A4',
                '#0E2DEC',
                '#6465CB',
                "#08A325",
                '#24F5A8',
                '#B80115']
    node_shape = ['s', 'o', '^', '>', 'v', '<', 'd', 'p', 'h', '8']

    for i in range(len(com)):
        nx.draw_networkx_nodes(G, pos,
                               nodelist=com[i],
                               node_color=color_list[i],
                               node_shape=node_shape[i])
    plt.show()


if __name__ == "__main__":
    # KL算法
    com = list(kernighan_lin_bisection(G))
    print('社区数量', len(com))
    for i in range(len(com)):
        print(com[i])
    draw_spring(G, com)
